# -*- coding: utf-8 -*-  
'''
baidu关系抽取数据集

@author: luoyi
Created on 2021年6月22日
'''
import json

import utils.conf as conf


#    样本定义
class Sample:
    #    文本内容
    @property
    def text(self):
        '''原字符串
            例如：如何演好自己的角色，请读《演员自我修养》《喜剧之王》周星驰崛起于穷困潦倒之中的独门秘笈
        '''
        return self._text
    
    #    词性标注
    @property
    def postag(self):
        '''dict 列表
            每个dict: {word:词字符串, loc:(头指针, 尾指针), pos:词性}
            例如：{word:周星驰, loc:(26, 28), pos:nr}
        '''
        return self._postag
    
    #    关系列表
    @property
    def spo_list(self):
        '''dict 列表
            每个dict: {subject:{word:头实体字符串, loc:(头指针, 尾指针), type:类型},
                      predicate:关系名称,
                      object:{word:头实体字符串, loc:(头指针, 尾指针), type:类型}
            例如：{subject:{word:周星驰, loc:(26, 28), type:人物},
                  predicate:主演,
                  object:{word:喜剧之王, loc:(21, 24), type:影视作品}
        '''
        return self._spo_list
    
    def __init__(self, json_fragment):
        #    解析原字符串
        self._text = json_fragment['text']
        
        #    解析词性标注
        self._postag = []
        pos_dict = {}            #    所有被标注的词在原句中的位置，后面关系中的词会检测是否存在该词
        i = 0
        for postag in json_fragment['postag']:
            #    检测被标注词性的词在原句中是否找得到
            loc, word = Sample.loc_from_text(self._text, postag['word'], i)
            if (loc is None):
                print('序列标注中 关键词在原句中找不到，忽略这条实体. text:', self._text, ' word:', postag['word'])
                continue
            postag['loc'] = loc
            postag['word'] = word
            self._postag.append(postag)
            
            i = loc[1] + 1
            pos_dict[word] = loc 
            pass

        #    解析关系
        self._spo_list = []
        i = 0
        for spo in json_fragment['spo_list']:
            #    检测subject在标注中是否能找到
            word_s = spo['subject']
            loc_s = pos_dict.get(word_s)
            if (loc_s is None):
                print('关系中 subject关键词在原句中找不到，忽略这条关系. text:', self._text, ' word:', spo['subject'])
                continue
            
            #    检测object在标注中是否能找到
            word_o = spo['object']
            loc_o = pos_dict.get(word_o)
            if (loc_o is None):
                print('关系中 object关键词在原句中找不到，忽略这条关系. text:', self._text, ' word:', spo['object'])
                continue
            
            self._spo_list.append({
                                    'subject': {'word':word_s, 'loc':loc_s, 'type':spo['subject_type']},
                                    'object': {'word':word_o, 'loc':loc_o, 'type':spo['object_type']},
                                    'predicate': spo['predicate']
                                 })
            pass
        pass
    
    #    找到词在原文本中的头尾指针
    @classmethod
    def loc_from_text(cls, text, word, i):
        #    通过首指针根据词长度从原句中取，如果取出来跟词不一致，则判定为标注错误
        word_ = text[i : i+len(word)]
        if (word != word_): return None, None
        #    当前指针即为首指针
        loc_s = i
        #    当前指针+词长度-1即为尾指针。（索引是包含尾指针的）
        loc_e = loc_s + len(word) - 1
        if (loc_e >= len(text)): return None, None
        
        return (loc_s, loc_e), word
    pass


#    遍历json文件
def dataset_iterator(file_path=conf.DATASET_BAIDU.get_train_data_path()):
    fr = open(file=file_path, mode='r', encoding='utf-8')
    while True:
        line = fr.readline()
        #    如果读不到内容，说明读到结尾了
        if (not line): break
        #    如果读到的空行，跳过
        if (len(line) == 0 or len(line.strip()) == 0): continue
        
        #    解析为json
        j = json.loads(line)
        sample = Sample(json_fragment=j)
        
        #    剔除掉一些不合法的数据
        if (len(sample.spo_list) == 0
            or len(sample.postag) == 0): 
            print('解析出的关系或实体为空，忽略这条数据. text:', sample.text)
            continue
        
        yield sample
        pass
    fr.close()
    pass

